import cv2
import os
import numpy as np
from tqdm import tqdm
import multiprocessing
from multiprocessing import Pool
from functools import partial
import matplotlib.pyplot as plt

class EnhancedImageStabilizer:
    def __init__(self):
        self.reference_frame = None
        self.reference_right_edge = None
        self.crop_width = None
        self.debug_mode = False
        self.roi = None  # 新增：存储用户标注的感兴趣区域
        self.consecutive_failures = 0  # 新增：连续检测失败计数
        
    def set_debug_mode(self, debug_mode=True):
        """设置调试模式，会保存中间过程图像以便观察"""
        self.debug_mode = debug_mode
    
    def set_roi(self, x, y, w, h):
        """设置用户标注的感兴趣区域"""
        self.roi = (x, y, w, h)
        print(f"已设置ROI区域: x={x}, y={y}, width={w}, height={h}")
        # 重置连续失败计数
        self.consecutive_failures = 0
    
    def detect_vertical_lines(self, gray_img):
        """检测图像中的垂直线条，改进版本"""
        # 如果设置了ROI，只处理ROI区域
        if self.roi is not None:
            x, y, w, h = self.roi
            roi_img = gray_img[y:y+h, x:x+w]
        else:
            roi_img = gray_img
        
        # 增强对比度
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
        enhanced = clahe.apply(roi_img)
        
        # 使用Canny边缘检测
        edges = cv2.Canny(enhanced, 50, 150)
        
        # 使用形态学操作增强垂直线条
        kernel_vertical = np.ones((15,1), np.uint8)
        vertical_edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel_vertical)
        vertical_edges = cv2.morphologyEx(vertical_edges, cv2.MORPH_OPEN, kernel_vertical)
        
        # 保存中间结果用于调试
        if self.debug_mode:
            cv2.imwrite("debug_enhanced.jpg", enhanced)
            cv2.imwrite("debug_edges.jpg", edges)
            cv2.imwrite("debug_vertical_edges.jpg", vertical_edges)
        
        return vertical_edges
    
    def find_right_edge(self, binary_img):
        """在二值化图像中找到右边缘，改进版本"""
        # 计算垂直投影
        projection = np.sum(binary_img, axis=0)
        
        # 使用一维高斯卷积减少噪声影响
        window = 15
        sigma = window/3.0
        x = np.linspace(-window//2, window//2, window)
        gaussian = np.exp(-(x**2)/(2*sigma**2))
        gaussian = gaussian/np.sum(gaussian)
        smoothed = np.convolve(projection, gaussian, mode='valid')
        
        # 计算梯度
        gradient = np.gradient(smoothed)
        
        # 使用更强的阈值检测边缘
        threshold = np.std(gradient) * 2.5
        
        # 从右向左扫描找到第一个显著边缘
        right_edge = None
        for i in range(len(gradient)-1, window, -1):
            if abs(gradient[i]) > threshold:
                # 使用局部最大值优化边缘位置
                local_region = gradient[max(0, i-8):min(len(gradient), i+9)]
                local_offset = np.argmax(np.abs(local_region))
                right_edge = i - 8 + local_offset
                break
        
        # 如果没有找到显著边缘，尝试使用更宽的搜索窗口和更低的阈值
        if right_edge is None:
            threshold = np.std(gradient) * 1.5
            for i in range(len(gradient)-1, window, -1):
                if abs(gradient[i]) > threshold:
                    local_region = gradient[max(0, i-10):min(len(gradient), i+11)]
                    local_offset = np.argmax(np.abs(local_region))
                    right_edge = i - 10 + local_offset
                    break
        
        # 如果设置了ROI，需要将检测到的边缘位置转换回原图坐标系
        if self.roi is not None and right_edge is not None:
            right_edge += self.roi[0]
        
        # 如果调试模式，保存检测结果
        if self.debug_mode and right_edge is not None:
            debug_img = cv2.cvtColor(binary_img, cv2.COLOR_GRAY2BGR)
            cv2.line(debug_img, (right_edge, 0), (right_edge, binary_img.shape[0]), (0, 0, 255), 2)
            cv2.imwrite("debug_right_edge.jpg", debug_img)
            
            plt.figure(figsize=(16, 12))
            plt.subplot(311)
            plt.title("Binary Image")
            plt.imshow(binary_img, cmap='gray')
            
            plt.subplot(312)
            plt.title("Projection Profile")
            plt.plot(projection)
            plt.plot(smoothed, 'r')
            if right_edge is not None:
                plt.axvline(x=right_edge, color='r')
            
            plt.subplot(313)
            plt.title("Gradient")
            plt.plot(gradient)
            plt.axhline(y=threshold, color='r', linestyle='--')
            plt.axhline(y=-threshold, color='r', linestyle='--')
            if right_edge is not None:
                plt.axvline(x=right_edge, color='r')
            
            plt.tight_layout()
            plt.savefig("debug_right_edge_analysis.jpg")
            plt.close()
            
        return right_edge
    
    def initialize_with_reference(self, reference_img):
        """使用参考帧初始化稳定器"""
        self.reference_frame = reference_img.copy()
        
        # 显示参考帧，让用户标注感兴趣区域
        cv2.imshow("标注参考帧", reference_img)
        print("请在图像上标注棒束通道和胶布的交界区域（矩形框）")
        print("按下鼠标左键开始绘制，松开鼠标左键完成绘制")
        print("按's'键保存标注，按'q'键退出")
        
        # 等待用户标注
        self.roi = cv2.selectROI("标注参考帧", reference_img, False)
        cv2.destroyAllWindows()
        
        if self.roi == (0, 0, 0, 0):
            print("警告：未标注区域，将使用全图检测")
            self.roi = None
        else:
            print(f"已标注ROI区域: {self.roi}")
        
        # 转换为灰度图并增强对比度
        gray = cv2.cvtColor(reference_img, cv2.COLOR_BGR2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        
        # 检测垂直线
        binary = self.detect_vertical_lines(gray)
        
        # 找到右边缘
        right_edge = self.find_right_edge(binary)
        
        if right_edge is not None:
            self.reference_right_edge = right_edge
            # 设置裁剪宽度（从右边缘向左的固定距离）
            self.crop_width = right_edge - 100  # 可以根据需要调整这个值
            print(f"初始化完成：参考右边缘位置: {right_edge}")
            print(f"裁剪宽度设置为: {self.crop_width}")
            return True
        else:
            print("警告：无法在参考帧中检测到右边缘！")
            return False
    
    def stabilize_image(self, img):
        """稳定图像，通过检测右边缘并调整裁剪区域"""
        if self.reference_right_edge is None or self.crop_width is None:
            return img, False
        
        # 转换为灰度图并增强对比度
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        
        # 检测垂直线
        binary = self.detect_vertical_lines(gray)
        
        # 找到当前帧的右边缘
        current_right_edge = self.find_right_edge(binary)
        
        if current_right_edge is None:
            print("警告：当前帧未检测到右边缘，使用参考帧的裁剪参数")
            # 使用参考帧的裁剪参数
            start_x = max(0, self.reference_right_edge - self.crop_width)
            end_x = self.reference_right_edge
            cropped = img[:, start_x:end_x]
            return cropped, False
        else:
            # 计算裁剪区域，使用动态边缘位置
            start_x = max(0, current_right_edge - self.crop_width)
            end_x = current_right_edge
            cropped = img[:, start_x:end_x]
            return cropped, True

def process_single_image(args):
    img_path, cropper = args
    try:
        # 读取当前帧
        current_frame = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        if current_frame is None:
            return None, False
        
        # 裁剪图像并获取检测状态
        cropped_frame, detection_success = cropper.stabilize_image(current_frame)
        
        # 返回处理后的帧和检测状态
        return cropped_frame, detection_success
        
    except Exception as e:
        print(f"处理图片 {img_path} 时发生错误: {str(e)}")
        return None, False

def process_images_to_video(input_path, output_path, sample_count=4000, debug_mode=False):
    try:
        # 获取所有bmp文件
        image_files = [f for f in os.listdir(input_path) if f.lower().endswith('.bmp')]
        if not image_files:
            print(f"错误：在 {input_path} 中没有找到BMP文件")
            return
        image_files.sort()
        
        total_images = len(image_files)
        if total_images < sample_count:
            sample_count = total_images
        step = total_images // sample_count
        
        # 读取第一张图片作为参考帧
        first_image_path = os.path.join(input_path, image_files[0])
        print(f"正在读取第一张图片: {first_image_path}")
        reference_frame = cv2.imdecode(np.fromfile(first_image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        if reference_frame is None:
            print(f"错误：无法读取图片 {first_image_path}")
            return
        
        # 创建裁剪器并初始化
        cropper = EnhancedImageStabilizer()
        cropper.set_debug_mode(debug_mode)
        if not cropper.initialize_with_reference(reference_frame):
            return
        
        # 获取裁剪后的尺寸
        first_processed, _ = cropper.stabilize_image(reference_frame)
        height, width = first_processed.shape[:2]
        
        # 准备所有需要处理的文件
        selected_files = [os.path.join(input_path, image_files[i]) for i in range(0, total_images, step)]
        
        # 确保输出目录存在
        os.makedirs(output_path, exist_ok=True)
        
        # 生成视频文件路径
        video_path = os.path.join(output_path, "output_video_stabilized.mp4")
        print(f"正在生成视频: {video_path}")
        
        # 创建视频写入器
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(video_path, fourcc, 30.0, (width, height))
        
        if not out.isOpened():
            print(f"错误：无法创建视频文件 {video_path}")
            return
        
        # 处理所有图片
        consecutive_failures = 0
        batch_size = 100  # 每批处理的图片数量
        processed_frames = []
        total_files = len(selected_files)
        
        with tqdm(total=total_files, desc="总体进度") as pbar:
            for i in range(0, total_files, batch_size):
                batch_files = selected_files[i:i + batch_size]
                process_args = [(f, cropper) for f in batch_files]
                
                # 获取CPU核心数
                num_cores = multiprocessing.cpu_count()
                
                # 创建进程池处理当前批次
                with Pool(num_cores) as pool:
                    # 处理当前批次
                    batch_frames = list(pool.imap(process_single_image, process_args))
                
                # 更新总进度
                pbar.update(len(batch_files))
                
                # 统计检测失败的帧数
                failed_frames = sum(1 for frame, success in batch_frames if not success)
                print(f"\n当前批次检测失败帧数: {failed_frames}")  # 添加调试信息
                
                if failed_frames >= 3:
                    consecutive_failures += 1
                    print(f"连续失败批次数: {consecutive_failures}")  # 添加调试信息
                else:
                    consecutive_failures = 0
                
                # 如果连续三批次都有大量检测失败，暂停处理让用户重新标注
                if consecutive_failures >= 3:
                    print("\n连续多次检测失败，需要重新标注区域")
                    # 读取当前批次的最后一帧用于标注
                    current_frame = cv2.imdecode(
                        np.fromfile(batch_files[-1], dtype=np.uint8),
                        cv2.IMREAD_COLOR
                    )
                    
                    # 显示当前帧让用户重新标注
                    cv2.imshow("重新标注区域", current_frame)
                    print("请重新标注棒束通道和胶布的交界区域（矩形框）")
                    print("按下鼠标左键开始绘制，松开鼠标左键完成绘制")
                    print("按's'键保存标注，按'q'键退出")
                    
                    # 等待用户标注
                    new_roi = cv2.selectROI("重新标注区域", current_frame, False)
                    cv2.destroyAllWindows()
                    
                    if new_roi != (0, 0, 0, 0):
                        cropper.set_roi(*new_roi)
                        consecutive_failures = 0
                        # 重新处理当前批次
                        continue
                
                # 将成功处理的帧添加到结果列表
                processed_frames.extend([frame for frame, success in batch_frames if frame is not None])
        
        if not processed_frames:
            print("错误：没有成功处理任何图片")
            return
        
        # 写入处理后的帧
        print("\n正在写入视频...")
        for frame in processed_frames:
            out.write(frame)
        
        out.release()
        print(f"视频生成完成！保存在: {video_path}")
        
    except Exception as e:
        print(f"发生错误: {str(e)}")

def preview_stabilization(input_path, output_dir, debug_mode=True):
    """预览稳定效果，处理少量图片并保存结果"""
    try:
        # 获取所有bmp文件
        image_files = [f for f in os.listdir(input_path) if f.lower().endswith('.bmp')]
        if not image_files:
            print(f"错误：在 {input_path} 中没有找到BMP文件")
            return
        image_files.sort()
        
        # 选择少量图片用于预览
        total_images = len(image_files)
        sample_count = min(100, total_images)  # 最多10张图片用于预览
        step = max(1, total_images // sample_count)
        
        # 读取第一张图片作为参考帧
        first_image_path = os.path.join(input_path, image_files[0])
        print(f"正在读取第一张图片: {first_image_path}")
        reference_frame = cv2.imdecode(np.fromfile(first_image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        if reference_frame is None:
            print(f"错误：无法读取图片 {first_image_path}")
            return
        
        # 创建裁剪器并初始化
        cropper = EnhancedImageStabilizer()
        cropper.set_debug_mode(debug_mode)
        if not cropper.initialize_with_reference(reference_frame):
            return
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 处理选定的图片
        selected_files = [image_files[i] for i in range(0, total_images, step)]
        
        for i, filename in enumerate(selected_files):
            try:
                # 读取当前帧
                img_path = os.path.join(input_path, filename)
                current_frame = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
                
                if current_frame is None:
                    print(f"警告：无法读取图片 {img_path}")
                    continue
                
                # 保存原始图像
                original_output_path = os.path.join(output_dir, f"original_{i+1}.jpg")
                cv2.imwrite(original_output_path, current_frame)
                
                # 稳定图像
                stabilized_frame, _ = cropper.stabilize_image(current_frame)
                
                # 保存稳定后的图像
                stabilized_output_path = os.path.join(output_dir, f"stabilized_{i+1}.jpg")
                cv2.imwrite(stabilized_output_path, stabilized_frame)
                
                print(f"处理完成 {i+1}/{len(selected_files)}: {filename}")
                
            except Exception as e:
                print(f"处理图片 {filename} 时发生错误: {str(e)}")
        
        print("预览图片生成完成！")
        
    except Exception as e:
        print(f"发生错误: {str(e)}")

if __name__ == '__main__':
    # 使用原始字符串表示路径
    input_path = r"H:\李双成\E-1000\C001H001S0001"  # 原始图像路径
    output_path = r"H:\李双成\test"  # 输出视频目录
    preview_dir = r"H:\test\preview_stabilization"  # 预览图像保存路径
    
    # 先预览效果
    preview_stabilization(input_path, preview_dir, debug_mode=True)
    
    # 根据预览结果确认算法效果后再处理所有图片生成视频
    # process_images_to_video(input_path, output_path, debug_mode=False)
